import matplotlib.pyplot as plt  # 导入matplotlib.pyplot模块，用于绘图
import numpy as np  # 导入numpy模块，用于数值计算
from gzip import open as gzip_open  # 从gzip模块导入open函数，用于读取gzip文件
from pickle import load as pickle_load  # 从pickle模块导入load函数，用于加载pickle文件
from pathlib import Path  # 从pathlib模块导入Path类，用于处理文件路径

from src.config.config import Config  # 从src.config.config模块导入Config类，用于获取配置
from src.config.config_plotting import PlotConfig, save_figures, generic_styling  # 从src.config.config_plotting模块导入PlotConfig类、save_figures函数和generic_styling函数

def plot_distance_sweep_testing_graph(
        paths,  # 数据文件路径列表
        name,  # 图形名称
        width,  # 图形宽度
        height,  # 图形高度
        plots_parent_path,  # 图形保存路径
        legend: list or None = None,  # 图例标签，默认为None
        colors: list or None = None,  # 图形颜色，默认为None
        markerstyle: list or None = None,  # 标记样式，默认为None
        linestyles: list or None = None,  # 线型，默认为None
) -> None:

    # 创建一个绘图窗口和坐标轴
    fig, ax = plt.subplots(figsize=(width, height))

    data = []
    # 读取数据文件
    for path in paths:
        with gzip_open(path, 'rb') as file:  # 打开gzip文件
            data.append(pickle_load(file))  # 使用pickle加载数据，并添加到数据列表中
    
    # 遍历每个数据条目并绘制图形
    for data_id, data_entry in enumerate(data):
        first_entry = list(data_entry[1]['sum_rate'].keys())[0]  # 获取第一个数据条目的键
        
        # 设置绘图样式
        marker = markerstyle[data_id] if markerstyle is not None else None
        color = colors[data_id] if colors is not None else None
        linestyle = linestyles[data_id] if linestyles is not None else None

        # 设置markevery的值
        if len(data_entry[0][data_entry[1]['sum_rate'][first_entry]['mean'] > 0.999 * np.max(data_entry[1]['sum_rate'][first_entry]['mean'])]) < int(len(data_entry[0])/10):
            markevery = np.searchsorted(data_entry[0], data_entry[0][data_entry[1]['sum_rate'][first_entry]['mean'] > 0.999 * np.max(data_entry[1]['sum_rate'][first_entry]['mean'])])
            for value_id in reversed(range(1, len(markevery))):
                if markevery[value_id] / markevery[value_id-1] < 1.01:
                    markevery = np.delete(markevery, value_id)
        else:
            markevery = (int(len(data_entry[0])/10 / len(data)*data_id), int(len(data_entry[0])/10))

        # 绘制数据曲线
        ax.plot(data_entry[0],  # X轴数据
                data_entry[1]['sum_rate'][first_entry]['mean'],  # Y轴数据
                color=color,  # 线条颜色
                marker=marker,  # 标记样式
                markevery=markevery,  # 每隔多少个数据点绘制标记
                linestyle=linestyle)  # 线型

    # 设置坐标轴标签
    ax.set_xlabel('User Distance \( D_{usr} \) [m]')  # X轴标签
    ax.set_ylabel('Sum Rate \( R \) [bps/Hz]')  # Y轴标签

    # 添加图例
    if legend:
        ax.legend(legend, ncols=3)  # 设置图例并分成3列

    # 应用通用样式
    generic_styling(ax=ax)
    fig.tight_layout(pad=0)  # 自动调整布局

    # 保存图形
    save_figures(plots_parent_path=plots_parent_path, plot_name=name, padding=0)


if __name__ == '__main__':  # 确保仅在直接运行脚本时执行以下代码

    cfg = Config()  # 创建Config类的实例，用于加载配置信息
    plot_cfg = PlotConfig()  # 创建PlotConfig类的实例，用于加载绘图配置信息

    # 定义数据文件路径列表
    data_paths = [
        Path(cfg.output_metrics_path,
             'sat_2_ant_4_usr_3_satdist_10000_usrdist_1000', 'distance_sweep',
             'testing_mmse_sweep_970.0_1029.9899999999454.gzip'),

        Path(cfg.output_metrics_path,
             'sat_2_ant_4_usr_3_satdist_10000_usrdist_1000', 'distance_sweep',
             'testing_sac_error_0.0_userwiggle_30_snap_3.580_sweep_970.0_1029.9899999999454.gzip'),
        # Path(cfg.output_metrics_path,
        #      'sat_2_ant_4_usr_3_satdist_10000_usrdist_1000', 'distance_sweep',
        #      'testing_sac_error_0.1_userwiggle_30_snap_3.422_sweep_970.0_1029.9899999999454.gzip'),
        Path(cfg.output_metrics_path,
             'sat_2_ant_4_usr_3_satdist_10000_usrdist_1000', 'distance_sweep',
             'testing_mrc_sweep_970.0_1029.9899999999454.gzip'),
    ]

    # 设置图形的宽度和高度
    plot_width = 0.99 * plot_cfg.textwidth  # 图形宽度为plot_cfg.textwidth的99%
    plot_height = plot_width * 1/3  # 图形高度为宽度的三分之一

    # 定义图例标签、标记样式、颜色和线型
    plot_legend = ['MMSE', 'SAC1', 'OMA']  # 图例标签
    plot_markerstyle = ['o', 's', '^']  # 标记样式，分别为圆圈、方形和三角形
    plot_colors = [plot_cfg.cp2['blue'], plot_cfg.cp2['magenta'], plot_cfg.cp2['black']]  # 图形颜色，分别为蓝色、洋红色和黑色
    plot_linestyles = ['-', '-', '--']  # 线型，实线和虚线

    # 调用自定义绘图函数，生成距离扫描测试图
    plot_distance_sweep_testing_graph(
        paths=data_paths,  # 数据文件路径
        name='dist_sweep_test_long',  # 图形名称
        width=plot_width,  # 图形宽度
        height=plot_height,  # 图形高度
        legend=plot_legend,  # 图例标签
        colors=plot_colors,  # 图形颜色
        markerstyle=plot_markerstyle,  # 标记样式
        linestyles=plot_linestyles,  # 线型
        plots_parent_path=plot_cfg.plots_parent_path,  # 图形保存路径
    )
    
    plt.show()  # 显示绘制的图形
